# ==============================================================================
# Copyright (c) 2025 CompanyName. All rights reserved.
# Author:         22020873 陈泽欣
# Project:        Design of Deep Learning Fundamental Course
# Module:         pose_estimationGUI.py
# Date:           2025-05-24
# Description:    本模块实现了骰子姿态识别系统的图形用户界面（GUI），提供友好的交互方式以支持图像加载、实时检测、结果可视化与参数调节。
#                 主要功能包括：
#                 - 多种输入源选择（图片、视频、屏幕截图）；
#                 - 实时显示处理结果与多种可视化选项；
#                 - 支持在线调整关键算法参数（如置信度阈值、重投影误差阈值等）；
#                 是整个系统中人机交互与流程控制的核心界面模块。
# ==============================================================================

import sys
from PyQt5.QtWidgets import (QApplication, QMainWindow, QLabel, QPushButton,
                             QVBoxLayout, QHBoxLayout, QGroupBox, QRadioButton,
                             QCheckBox, QWidget, QFileDialog, QMessageBox,
                             QLineEdit, QFormLayout, QGridLayout
                             )
from PyQt5.QtGui import QImage, QPixmap, QIcon
from PyQt5.QtCore import Qt, QTimer
import cv2
from core.get_features import process_frame
from utils.grabscreen import grab_screen, setup_window
from configs.config import Config
from configs.norm_data import btn_start_styleSheet, btn_stop_styleSheet, line_edit_styleSheet


def set_checkbox_style(check_box, checked_color_bgr = (0, 191, 255), unchecked_color="#FFFFFF"):
    check_box.setStyleSheet(f"""
        QCheckBox {{
            color: black;
            font-size: 13px;
            margin-left: 5px;
        }}
        QCheckBox::indicator {{
            width: 10px;
            height: 10px;
            border: 1px solid #DCDCDC; 
            border-radius: 6px;
            background-color: {unchecked_color}; 
        }}
        QCheckBox::indicator:checked {{
            background-color: rgb({checked_color_bgr[2]}, {checked_color_bgr[1]}, {checked_color_bgr[0]});
            image: url(path/to/checked_icon.png); 
        }}
        QCheckBox::indicator:unchecked {{
            background-color: {unchecked_color}; 
            image: url(path/to/unchecked_icon.png); 
        }}
    """)

class PoseEstimationUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle(Config.WINDOW_TITLE)
        self.resize(Config.WINDOW_WIDTH, Config.WINDOW_HEIGHT)

        self.setWindowFlags(self.windowFlags() | Qt.WindowStaysOnTopHint)

        # 状态变量
        self.image_source = "image"
        self.cap = None
        self.timer = QTimer()
        self.current_img = None

        self.setWindowIcon(QIcon('gui\\icon.ico'))

        # 初始化 UI
        self.init_ui()

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout()

        # ========================
        # 图像显示区域（最上方）
        # ========================
        self.image_label = QLabel("图像预览区")
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setStyleSheet("border: 2px solid gray;")
        self.image_label.setFixedSize(Config.WINDOW_WIDTH, Config.WINDOW_HEIGHT)
        main_layout.addWidget(self.image_label, alignment=Qt.AlignCenter)

        # ========================
        # 控制面板（水平布局：输入源 + 显示选项 + 参数调节）
        # ========================
        control_layout = QHBoxLayout()

        # 左侧：输入源选择（占 1/4）
        input_group = QGroupBox("输入源")
        param_grid = QGridLayout()
        input_layout = QVBoxLayout()

        self.radio_image = QRadioButton("加载图片")
        self.radio_video = QRadioButton("加载视频")
        self.radio_realtime = QRadioButton("实时检测")
        self.radio_image.setChecked(True)

        self.radio_image.toggled.connect(lambda: self.set_input_source("image"))
        self.radio_video.toggled.connect(lambda: self.set_input_source("video"))
        self.radio_realtime.toggled.connect(lambda: self.set_input_source("realtime"))

        self.btn_start = QPushButton("Start")
        self.btn_start.clicked.connect(self.load_file)

        self.btn_stop = QPushButton("Stop")
        self.btn_stop.clicked.connect(self.stop_current_source)

        self.btn_start.setStyleSheet(btn_start_styleSheet)
        self.btn_stop.setStyleSheet(btn_stop_styleSheet)

        input_layout.addWidget(self.radio_image)
        input_layout.addWidget(self.radio_video)
        input_layout.addWidget(self.radio_realtime)
        input_layout.addWidget(self.btn_start)
        input_layout.addWidget(self.btn_stop)
        input_group.setLayout(input_layout)
        control_layout.addWidget(input_group, stretch=1)

        # 中间：结果显示选项（占 1/4）
        display_group = QGroupBox("显示选项")
        display_layout = QVBoxLayout()
        self.checkbox_show_yolo = QCheckBox("展示YOLO结果")
        self.checkbox_pnp = QCheckBox("PnP 结果")
        self.checkbox_ba = QCheckBox("BA 结果")
        self.checkbox_posenet = QCheckBox("PoseNet 结果")
        self.checkbox_fusion = QCheckBox("融合结果")

        set_checkbox_style(self.checkbox_show_yolo, (50, 220, 255))
        set_checkbox_style(self.checkbox_pnp, Config.PNP_FRAME_COLOR)
        set_checkbox_style(self.checkbox_ba, Config.BA_FRAME_COLOR)
        set_checkbox_style(self.checkbox_posenet, Config.POSENET_FRAME_COLOR)
        set_checkbox_style(self.checkbox_fusion, Config.FUSION_FRAME_COLOR)

        self.checkbox_show_yolo.setChecked(True)
        self.checkbox_pnp.setChecked(True)
        self.checkbox_fusion.setChecked(True)

        self.checkbox_show_yolo.stateChanged.connect(self.toggle_show_yolo)
        self.checkbox_pnp.stateChanged.connect(self.toggle_pnp)
        self.checkbox_ba.stateChanged.connect(self.toggle_ba)
        self.checkbox_posenet.stateChanged.connect(self.toggle_posenet)
        self.checkbox_fusion.stateChanged.connect(self.toggle_fusion)

        display_layout.addWidget(self.checkbox_show_yolo)
        display_layout.addWidget(self.checkbox_pnp)
        display_layout.addWidget(self.checkbox_ba)
        display_layout.addWidget(self.checkbox_posenet)
        display_layout.addWidget(self.checkbox_fusion)
        display_group.setLayout(display_layout)
        control_layout.addWidget(display_group, stretch=1)

        # 右侧：参数调节（占 1/2）
        param_group = QGroupBox("参数调节")
        param_layout = QFormLayout()

        self.param_inputs = {}

        default_params = {
            "yolo_conf": Config.ALLOW_YOLO_CONFIDENCE_THRESHOLD,
            "normal_repro_err": Config.NORMAL_MIN_ALLOW_REPROJECTION_ERROR,
            "posenet_repro_err": Config.POSENET_MIN_ALLOW_REPROJECTION_ERROR,
            "fusion_repro_err": Config.FUSION_MIN_ALLOW_REPROJECTION_ERROR,
            "contour_area": Config.CONTOURAREA_THRESHOLD,
            "huber_loss": Config.HUBER_LOSS_VALUE,
            "box_offset_value": Config.BOX_OFFSET_VALUE,
            "use_hsv_flag": Config.USE_HSV_FLAG,

            "normal_line_thickness": Config.NORMAL_LINE_THICKNESS,
            "text_font_thickness": Config.TEXT_FONT_THICKNESS,
            "text_font_scale": Config.TEXT_FONT_SCALE,
            "yolo_font_scale": Config.YOLO_FONT_SCALE,
            "pnp_frame_color": Config.PNP_FRAME_COLOR,
            "ba_frame_color": Config.BA_FRAME_COLOR,
            "posenet_frame_color": Config.POSENET_FRAME_COLOR,
            "fusion_frame_color": Config.FUSION_FRAME_COLOR,
        }

        self.setStyleSheet(line_edit_styleSheet)

        row = 0
        col = 0
        for key, value in default_params.items():
            label = QLabel(key)
            line_edit = QLineEdit(str(value))
            self.param_inputs[key] = line_edit  # 存储 QLineEdit 控件本身

            # 添加回车事件监听器（按下 Enter 键触发 display_image）
            line_edit.returnPressed.connect(self.display_image)

            param_grid.addWidget(label, row, col * 2)  # 标签
            param_grid.addWidget(line_edit, row, col * 2 + 1)  # 输入框

            row += 1
            if row == 8:
                col += 1
                row = 0

        param_group.setLayout(param_grid)
        control_layout.addWidget(param_group, stretch=2)

        # 设置主布局
        main_layout.addLayout(control_layout)
        central_widget.setLayout(main_layout)

        # 启动默认输入源
        # self.start_default_source()

    # =====================================
    # 控制逻辑函数
    # =====================================
    def set_input_source(self, source):
        self.image_source = source

    def toggle_show_yolo(self, state):
        self.show_yolo_result = state == Qt.Checked

    def toggle_pnp(self, state):
        set_checkbox_style(self.checkbox_pnp, eval(self.param_inputs['pnp_frame_color'].text()))
        self.show_pnp = state == Qt.Checked

    def toggle_ba(self, state):
        set_checkbox_style(self.checkbox_ba, eval(self.param_inputs['ba_frame_color'].text()))
        self.show_ba = state == Qt.Checked

    def toggle_posenet(self, state):
        set_checkbox_style(self.checkbox_posenet, eval(self.param_inputs['posenet_frame_color'].text()))
        self.show_posenet = state == Qt.Checked

    def toggle_fusion(self, state):
        set_checkbox_style(self.checkbox_fusion, eval(self.param_inputs['fusion_frame_color'].text()))
        self.show_fusion = state == Qt.Checked

    def start_default_source(self):
        if self.image_source == "realtime":
            self.start_realtime()
        elif self.image_source == "image":
            self.load_file()
        elif self.image_source == "video":
            self.load_file()

    def load_file(self):
        if self.image_source == "image":
            file_name, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.png *.jpg *.bmp)")
            if file_name:
                self.current_img = cv2.imread(file_name)
                # self.current_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
                self.display_image()
        elif self.image_source == "video":
            file_name, _ = QFileDialog.getOpenFileName(self, "选择视频", "", "Videos (*.mp4 *.avi)")
            if file_name:
                self.play_video(file_name)
        elif self.image_source == "realtime":
            self.start_realtime()

    def display_image(self):
        if self.current_img is not None:
            display_options = {
                'show_yolo': self.checkbox_show_yolo.isChecked(),
                'show_pnp': self.checkbox_pnp.isChecked(),
                'show_ba': self.checkbox_ba.isChecked(),
                'show_posenet': self.checkbox_posenet.isChecked(),
                'show_fusion': self.checkbox_fusion.isChecked(),
            }

            # 获取所有 QLineEdit 的值并转换为合适的类型
            params = {}
            for key, le in self.param_inputs.items():
                text = le.text().strip()
                try:
                    if '.' in text:
                        params[key] = float(text)
                    elif '(' and ')' in text:
                        params[key] = eval(text)
                    elif len(text) == 0:
                        params[key] = ''
                    else:
                        params[key] = int(text)

                except ValueError:
                    params[key] = text  # 如果不是数字则保留字符串

            draw_img = process_frame(self.current_img, display_options, params)
            self.display_cv_image(draw_img)

    def play_video(self, video_path):
        self.cap = cv2.VideoCapture(video_path)
        if not self.cap.isOpened():
            QMessageBox.critical(self, "错误", "无法打开视频文件")
            return

        self.timer.timeout.connect(self.update_video_frame)
        self.timer.start(30)

    def start_realtime(self):
        self.timer.timeout.connect(self.update_realtime_frame)
        self.timer.start(30)

    def update_video_frame(self):
        ret, frame = self.cap.read()
        if ret:
            display_options = {
                'show_yolo': self.checkbox_show_yolo.isChecked(),
                'show_pnp': self.checkbox_pnp.isChecked(),
                'show_ba': self.checkbox_ba.isChecked(),
                'show_posenet': self.checkbox_posenet.isChecked(),
                'show_fusion': self.checkbox_fusion.isChecked(),
            }

            params = {}
            for key, le in self.param_inputs.items():
                text = le.text().strip()
                try:
                    if '.' in text:
                        params[key] = float(text)
                    elif '(' and ')' in text:
                        params[key] = eval(text)
                    elif len(text) == 0:
                        params[key] = ''
                    else:
                        params[key] = int(text)

                except ValueError:
                    params[key] = text

            draw_img = process_frame(frame, display_options, params)
            self.display_cv_image(draw_img)

    def update_realtime_frame(self):
        frame = grab_screen(region=(0, 0, Config.IMG_WIDTH, Config.IMG_HEIGHT))

        display_options = {
            'show_yolo': self.checkbox_show_yolo.isChecked(),
            'show_pnp': self.checkbox_pnp.isChecked(),
            'show_ba': self.checkbox_ba.isChecked(),
            'show_posenet': self.checkbox_posenet.isChecked(),
            'show_fusion': self.checkbox_fusion.isChecked(),
        }

        params = {}
        for key, le in self.param_inputs.items():
            text = le.text().strip()
            try:
                if '.' in text:
                    params[key] = float(text)
                elif '(' and ')' in text:
                    params[key] = eval(text)
                elif len(text) == 0:
                    params[key] = ''
                else:
                    params[key] = int(text)

            except ValueError:
                params[key] = text

        draw_img = process_frame(frame, display_options, params)
        self.display_cv_image(draw_img)

    def display_cv_image(self, cv_img):
        """将 OpenCV 图像转为 QImage 显示到 QLabel"""
        if cv_img is None:
            return

        h, w, ch = cv_img.shape
        bytes_per_line = ch * w
        q_img = QImage(cv_img.data, w, h, bytes_per_line, QImage.Format_BGR888)
        pixmap = QPixmap.fromImage(q_img).scaled(
            self.image_label.width(), self.image_label.height(), Qt.KeepAspectRatio
        )
        self.image_label.setPixmap(pixmap)

    def stop_current_source(self):
        if self.cap:
            self.cap.release()
            self.cap = None
        self.timer.stop()

        if self.image_source == "realtime":
            pass

        # 清除显示区域
        self.image_label.clear()

    def closeEvent(self, event):
        """关闭窗口时释放资源"""
        if self.cap:
            self.cap.release()
        self.timer.stop()
        event.accept()

